package com.xlf.service.ai.controller;

import cn.hutool.json.JSON;
import cn.hutool.json.JSONUtil;
import com.xlf.common.auth.util.JwtTokenService;
import com.xlf.common.auth.util.SecurityUtils;
import com.xlf.common.core.annotation.Anonymous;
import com.xlf.common.core.exception.BizException;
import com.xlf.common.core.pojo.mode.AjaxResult;
import com.xlf.common.core.pojo.mode.BaseController;
import com.xlf.service.ai.pojo.dto.AiRequest;
import com.xlf.service.ai.pojo.dto.AiResponse;
import com.xlf.service.ai.pojo.dto.ChatResponse;
import com.xlf.service.ai.service.AiService;
import com.xlf.service.ai.service.impl.ErnieBot4Client;
import com.xlf.service.ai.utils.SseEmitterUtil;
import lombok.SneakyThrows;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.http.MediaType;
import org.springframework.web.bind.annotation.*;
import org.springframework.web.servlet.mvc.method.annotation.SseEmitter;

import java.io.IOException;
import java.util.ArrayList;
import java.util.function.Consumer;
import java.util.stream.Collectors;
import java.util.stream.Stream;

/**
 * @Author xlf 2734029211@qq.com
 * @Date 2024/1/28 15:35
 * @Description: TODO
 * @Version 1.0
 */
@RestController
@RequestMapping("/ai")
public class AiController extends BaseController {

    @Autowired
    AiService aiService;
    @Autowired
    ErnieBot4Client ernieBot4Client;


    @PostMapping(value = "/chats")
    public AjaxResult chatSingle(@RequestBody AiRequest aiRequest) {
      return AjaxResult.success("生成成功",aiService.doAnswer(aiRequest));
    }


    /**
     * 建立sse连接
     * @return
     */
    @GetMapping(value = "/stream/connect",produces = {MediaType.TEXT_EVENT_STREAM_VALUE})
    public SseEmitter connect() {
        Long userId = (Long)SecurityUtils.getUserId();
        return SseEmitterUtil.connect(userId.toString());
    }


    @PostMapping(value = "/stream/chats")
    public AjaxResult chatSingleOfStream(@RequestBody AiRequest aiRequest) {

        Long userId = (Long) SecurityUtils.getUserId();

        if(!SseEmitterUtil.getIds().contains(userId.toString()))
            return error("未建立连接或连接已超时，请重连！");

        try {
            ernieBot4Client.chatContOfStream(aiRequest.getContent(),userId.toString() + ':' + aiRequest.getChatId(), new Consumer<ChatResponse>() {
                @SneakyThrows
                @Override
                public void accept(ChatResponse chatResponse) {
                    AiResponse aiResponse = new AiResponse();
                    aiResponse.setChatId(aiRequest.getChatId());
                    aiResponse.setResult(chatResponse.getResult());
                    aiResponse.setIsEnd(chatResponse.getIsEnd().equals(true) ? 1 : 0);
                    SseEmitterUtil.sendMessage(userId.toString(), JSONUtil.toJsonStr(aiResponse));
                }
            });
        } catch (Exception e) {
            return error("AI请求失败");
        } finally {
            SseEmitterUtil.removeUser(userId.toString());
        }

        return success("AI请求已完成");
    }


}
